#!/usr/bin/env python3
import time
import subprocess
import os

def configure_switch_rules():
    """Configure all 6 switches for bidirectional 4-path ECMP load balancing"""

    print("Configuring 6-switch bidirectional ECMP topology with 4 paths...")

    # s1: ECMP Load Balancer (connects to h1 and 4 paths)
    s1_commands = [
        # Forward ECMP group for virtual IP 10.0.0.1 with 4 paths (h1→h2)
        "table_add MyIngress.ecmp_group_forward MyIngress.set_ecmp_select 10.0.0.1/32 => 0 4",

        # Forward ECMP next hops - 4 paths through s2, s3, s4, s5 (FIXED: keep dst IP as 10.0.0.1)
        "table_add MyIngress.ecmp_nhop_forward MyIngress.set_nhop 0 => 00:00:00:00:02:00 10.0.0.1 2",  # to s2
        "table_add MyIngress.ecmp_nhop_forward MyIngress.set_nhop 1 => 00:00:00:00:03:00 10.0.0.1 3",  # to s3
        "table_add MyIngress.ecmp_nhop_forward MyIngress.set_nhop 2 => 00:00:00:00:04:00 10.0.0.1 4",  # to s4
        "table_add MyIngress.ecmp_nhop_forward MyIngress.set_nhop 3 => 00:00:00:00:05:00 10.0.0.1 5",  # to s5

        # Regular forwarding rules
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 08:00:00:00:01:01 1",  # to h1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 08:00:00:00:02:02 2",  # to h2 via any path
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.2/32 => 08:00:00:00:01:01 1",  # virtual IP to h1
    ]

    # s2: Path 1 switch
    s2_commands = [
        # Regular forwarding rules for path 1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 00:00:00:00:01:00 1",  # to s1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 00:00:00:00:06:00 2",  # to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.1/32 => 00:00:00:00:06:00 2",  # virtual IP to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.2/32 => 00:00:00:00:01:00 1",  # virtual IP to s1
    ]

    # s3: Path 2 switch
    s3_commands = [
        # Regular forwarding rules for path 2
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 00:00:00:00:01:00 1",  # to s1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 00:00:00:00:06:00 2",  # to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.1/32 => 00:00:00:00:06:00 2",  # virtual IP to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.2/32 => 00:00:00:00:01:00 1",  # virtual IP to s1
    ]

    # s4: Path 3 switch
    s4_commands = [
        # Regular forwarding rules for path 3
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 00:00:00:00:01:00 1",  # to s1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 00:00:00:00:06:00 2",  # to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.1/32 => 00:00:00:00:06:00 2",  # virtual IP to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.2/32 => 00:00:00:00:01:00 1",  # virtual IP to s1
    ]

    # s5: Path 4 switch
    s5_commands = [
        # Regular forwarding rules for path 4
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 00:00:00:00:01:00 1",  # to s1
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 00:00:00:00:06:00 2",  # to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.1/32 => 00:00:00:00:06:00 2",  # virtual IP to s6
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.2/32 => 00:00:00:00:01:00 1",  # virtual IP to s1
    ]

    # s6: Final hop to h2 with reverse ECMP
    s6_commands = [
        # Reverse ECMP group for virtual IP 10.0.0.2 with 4 paths (h2→h1)
        "table_add MyIngress.ecmp_group_reverse MyIngress.set_ecmp_select 10.0.0.2/32 => 0 4",

        # Reverse ECMP next hops - 4 paths through s2, s3, s4, s5 (FIXED: keep dst IP as 10.0.1.1)
        "table_add MyIngress.ecmp_nhop_reverse MyIngress.set_nhop 0 => 00:00:00:00:02:00 10.0.1.1 2",  # to s2
        "table_add MyIngress.ecmp_nhop_reverse MyIngress.set_nhop 1 => 00:00:00:00:03:00 10.0.1.1 3",  # to s3
        "table_add MyIngress.ecmp_nhop_reverse MyIngress.set_nhop 2 => 00:00:00:00:04:00 10.0.1.1 4",  # to s4
        "table_add MyIngress.ecmp_nhop_reverse MyIngress.set_nhop 3 => 00:00:00:00:05:00 10.0.1.1 5",  # to s5

        # Regular forwarding rules
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.1.1/32 => 08:00:00:00:01:01 2",  # to h1 via any path
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.2.2/32 => 08:00:00:00:02:02 1",  # to h2
        "table_add MyIngress.ipv4_lpm MyIngress.ipv4_forward 10.0.0.1/32 => 08:00:00:00:02:02 1",  # virtual IP to h2
    ]

    # Configure all 6 switches
    switch_configs = [
        ("s1", 9091, s1_commands),
        ("s2", 9092, s2_commands),
        ("s3", 9093, s3_commands),
        ("s4", 9094, s4_commands),
        ("s5", 9095, s5_commands),
        ("s6", 9096, s6_commands)
    ]

    for switch_name, thrift_port, commands in switch_configs:
        print(f"Configuring {switch_name}...")

        for cmd in commands:
            full_cmd = f'echo "{cmd}" | simple_switch_CLI --thrift-port {thrift_port}'
            try:
                result = subprocess.run(full_cmd, shell=True, capture_output=True, text=True)
                if result.returncode == 0:
                    if "Adding entry to table" in result.stdout or "Entry added" in result.stdout:
                        print(f"  ✅ {cmd}")
                    else:
                        print(f"  ⚠️  {cmd} (no confirmation)")
                else:
                    print(f"  ❌ {cmd}")
                    print(f"     Error: {result.stderr}")
            except Exception as e:
                print(f"  ❌ {cmd} - Exception: {e}")

def main():
    print("🚀 Starting 6-switch bidirectional ECMP configuration")
    print("Waiting for switches to initialize...")
    time.sleep(10)  # Give switches time to start

    configure_switch_rules()

    print("\n" + "="*60)
    print("🎯 Bidirectional ECMP Load Balancing Configuration Complete!")
    print("="*60)
    print("Topology: h1 ↔ s1 ↔ [s2,s3,s4,s5] ↔ s6 ↔ h2")
    print("Forward ECMP: 10.0.0.1 → Load balanced across 4 paths (h1→h2)")
    print("Reverse ECMP: 10.0.0.2 → Load balanced across 4 paths (h2→h1)")
    print("\n📊 To monitor ECMP:")
    print("   On s1: sudo tcpdump -i veth3 -c 10 -nn -q &  # s1->s2")
    print("   On s1: sudo tcpdump -i veth5 -c 10 -nn -q &  # s1->s3")
    print("   On s1: sudo tcpdump -i veth7 -c 10 -nn -q &  # s1->s4")
    print("   On s1: sudo tcpdump -i veth9 -c 10 -nn -q &  # s1->s5")
    print("\n🧪 Test commands:")
    print("   On h1: sudo ip netns exec h1 python3 -c \\\"from scapy.all import *; conf.L3socket = L3RawSocket; sendp(Ether(src='08:00:00:00:01:01', dst='08:00:00:00:01:00')/IP(src='10.0.1.1', dst='10.0.0.1')/TCP(dport=1234)/'Test', iface='eth0')\\\"")
    print("   On h2: sudo ip netns exec h2 ./receive.py")
    print("\nPress Ctrl+C to exit...")

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n🛑 Shutting down controller...")

if __name__ == '__main__':
    main()
